from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import shutil
import tempfile
import os
import sys
import uuid
from typing import Dict, List, Any, Optional

# 添加项目根目录到路径
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from core.processing_chain import DocumentGenerationChain
from tools.ollama_manager import OllamaManager
from tools.vector_store import VectorStoreManager
from data.loaders import CodeLoader
from data.splitters import CodeSplitter

app = FastAPI(title="代码文档生成API")

# 保存上传的项目
UPLOAD_DIR = "uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)

# 存储处理中的任务
processing_tasks = {}


class DocumentationRequest(BaseModel):
    """文档生成请求模型"""
    project_id: str
    model_name: str = "codellama"  # 默认使用codellama


class TaskStatus(BaseModel):
    """任务状态响应模型"""
    task_id: str
    status: str
    progress: float = 0
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None


@app.post("/upload_project/")
async def upload_project(file: UploadFile = File(...)):
    """上传项目ZIP文件

    Args:
        file: 上传的ZIP文件

    Returns:
        包含project_id的JSON响应
    """
    project_id = str(uuid.uuid4())
    project_dir = os.path.join(UPLOAD_DIR, project_id)
    os.makedirs(project_dir, exist_ok=True)

    # 保存上传的ZIP文件
    zip_path = os.path.join(project_dir, "project.zip")
    with open(zip_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    # 解压文件
    extracted_dir = os.path.join(project_dir, "extracted")
    os.makedirs(extracted_dir, exist_ok=True)

    try:
        shutil.unpack_archive(zip_path, extracted_dir)
    except Exception as e:
        return JSONResponse(
            status_code=400,
            content={"error": f"解压文件失败: {str(e)}"}
        )

    return {"project_id": project_id}


def process_documentation(project_id: str, model_name: str, task_id: str):
    """后台处理文档生成任务

    Args:
        project_id: 项目ID
        model_name: 模型名称
        task_id: 任务ID

    Returns:
        None
    """
    try:
        processing_tasks[task_id]["status"] = "processing"

        # 初始化组件
        ollama_manager = OllamaManager()
        llm = ollama_manager.get_model(model_name)

        vector_store_manager = VectorStoreManager()
        db = vector_store_manager.init_chroma(
            persist_directory=f"vectordb/{project_id}"
        )

        # 加载项目文件
        project_path = os.path.join(UPLOAD_DIR, project_id, "extracted")
        documents_by_file = CodeLoader.load_project(project_path)

        # 更新进度
        processing_tasks[task_id]["progress"] = 0.3

        # 分割文档
        split_documents = CodeSplitter.split_documents(documents_by_file)

        # 更新进度
        processing_tasks[task_id]["progress"] = 0.5

        # 添加到向量库
        all_splits = []
        for splits in split_documents.values():
            all_splits.extend(splits)

        vector_store_manager.add_documents(all_splits)

        # 更新进度
        processing_tasks[task_id]["progress"] = 0.7

        # 创建文档生成链
        doc_chain = DocumentGenerationChain(llm, db)
        doc_chain.setup_chain()

        # 生成文档
        original_docs = {}
        for file_path, docs in documents_by_file.items():
            if docs:  # 确保文档列表不为空
                original_docs[file_path] = docs[0].page_content

        # 处理所有文件并生成文档
        results = doc_chain.process_project(original_docs)

        # 更新进度和结果
        processing_tasks[task_id]["progress"] = 1.0
        processing_tasks[task_id]["status"] = "completed"
        processing_tasks[task_id]["result"] = results

    except Exception as e:
        processing_tasks[task_id]["status"] = "failed"
        processing_tasks[task_id]["error"] = str(e)
        print(f"处理任务出错: {e}")


@app.post("/generate_documentation/", response_model=TaskStatus)
async def generate_documentation(
        request: DocumentationRequest,
        background_tasks: BackgroundTasks
):
    """启动文档生成任务

    Args:
        request: 文档生成请求
        background_tasks: FastAPI后台任务

    Returns:
        任务状态
    """
    task_id = str(uuid.uuid4())

    # 检查项目是否存在
    project_dir = os.path.join(UPLOAD_DIR, request.project_id)
    if not os.path.exists(project_dir):
        raise HTTPException(status_code=404, detail="项目不存在")

    # 初始化任务状态
    processing_tasks[task_id] = {
        "status": "queued",
        "progress": 0,
        "result": None
    }

    # 启动后台任务
    background_tasks.add_task(
        process_documentation,
        request.project_id,
        request.model_name,
        task_id
    )

    return {"task_id": task_id, "status": "queued", "progress": 0}


@app.get("/task_status/{task_id}", response_model=TaskStatus)
async def get_task_status(task_id: str):
    """获取任务状态

    Args:
        task_id: 任务ID

    Returns:
        任务状态
    """
    if task_id not in processing_tasks:
        raise HTTPException(status_code=404, detail="任务不存在")

    task_info = processing_tasks[task_id]
    return {
        "task_id": task_id,
        "status": task_info["status"],
        "progress": task_info["progress"],
        "result": task_info.get("result"),
        "error": task_info.get("error")
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
